Skip to content

[ndarray] Add NDAxisIter struct and iter_by_axis method that iterate array by any axis #212

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
merged 3 commits into from
Feb 15, 2025
Merged

[ndarray] Add NDAxisIter struct and iter_by_axis method that iterate array by any axis #212

merged 3 commits into from
Feb 15, 2025

Conversation

@forfudan
Copy link
Collaborator

@forfudan forfudan commented Feb 13, 2025

Adds NDAxisIter struct and iter_by_axis method that iterate array by any axis. In each iteration, the iterator yields a 1-d array by specified axis. It is useful when we want to write a universal function to reduce the array by certain axis.

Example:

from numojo.prelude import *
var a = nm.arange[i8](24).reshape(Shape(2, 3, 4))
print(a)
for i in a.iter_by_axis(axis=0):
    print(String(i))

This prints:

[[[ 0  1  2  3]
    [ 4  5  6  7]
    [ 8  9 10 11]]
    [[12 13 14 15]
    [16 17 18 19]
    [20 21 22 23]]]
3D-array  Shape(2,3,4)  Strides(12,4,1)  DType: i8  C-cont: True  F-cont: False  own data: True
[ 0 12]
[ 1 13]
[ 2 14]
[ 3 15]
[ 4 16]
[ 5 17]
[ 6 18]
[ 7 19]
[ 8 20]
[ 9 21]
[10 22]
[11 23]

Another example:

from numojo.prelude import *
var a = nm.arange[i8](24).reshape(Shape(2, 3, 4))
print(a)
for i in a.iter_by_axis(axis=2):
    print(String(i))

This prints:

[[[ 0  1  2  3]
    [ 4  5  6  7]
    [ 8  9 10 11]]
    [[12 13 14 15]
    [16 17 18 19]
    [20 21 22 23]]]
3D-array  Shape(2,3,4)  Strides(12,4,1)  DType: i8  C-cont: True  F-cont: False  own data: True
[0 1 2 3]
[4 5 6 7]
[ 8  9 10 11]
[12 13 14 15]
[16 17 18 19]
[20 21 22 23]
```.

@shivasankarka shivasankarka merged commit 68b5595 into Mojo-Numerics-and-Algorithms-group:pre-0.6 Feb 15, 2025
2 checks passed
@forfudan forfudan deleted the iter branch February 15, 2025 11:15
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@shivasankarka shivasankarka shivasankarka approved these changes

@MadAlex1997 MadAlex1997 Awaiting requested review from MadAlex1997

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@forfudan @shivasankarka